# main.py
import os
import torch
import random
import numpy as np
import logging

from config.parser import parse_args
from data_loader.loader import DataLoader
from models.CS_KGAT import CS_KGAT
from trainer.trainer import Trainer
from utils.logger import logging_config, create_log_id


def main():
    """
    项目主入口函数，负责整个流程的协调。
    """
    # --- 1. 初始化与环境设置 ---
    args = parse_args()

    # 创建保存模型和日志的目录
    os.makedirs(args.save_dir, exist_ok=True)

    # 设置随机种子，确保实验可复现
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # 设置日志记录器
    log_save_id = create_log_id(args.save_dir)
    logging_config(folder=args.save_dir, name=f'log_{log_save_id}', no_console=False)
    logger = logging.getLogger()
    logger.info(f"Arguments: {args}")

    # 设置计算设备 (GPU or CPU)
    device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() and args.gpu_id != -1 else "cpu")
    logger.info(f"Using device: {device}")

    # --- 2. 加载和预处理数据 ---
    logger.info("--> Step 1: 开始加载和预处理数据...")
    try:
        data = DataLoader(args, logger)
    except Exception as e:
        logger.error(f"数据加载失败: {e}", exc_info=True)
        logger.error("请确保 'data/' 目录下已放置名为 '{}' 的数据集文件夹，并包含 train.txt, test.txt 等文件。".format(args.data_name))
        return
    logger.info("数据加载完成。")

    # --- 3. 【核心创新】加载LLM生成的语义嵌入 ---
    logger.info("--> Step 2: 加载LLM生成的语义画像嵌入...")
    user_profile_path = os.path.join(data.data_dir, 'user_profiles.pt')
    item_profile_path = os.path.join(data.data_dir, 'item_profiles.pt')

    try:
        user_profiles_embed = torch.load(user_profile_path, map_location='cpu')
        logger.info(f"用户画像加载成功: {user_profile_path}, Shape: {user_profiles_embed.shape}")

        item_profiles_embed = torch.load(item_profile_path, map_location='cpu')
        logger.info(f"物品画像加载成功: {item_profile_path}, Shape: {item_profiles_embed.shape}")
    except FileNotFoundError as e:
        # 【修正】: 如果文件不存在，直接报错，因为run.sh应该已经处理了生成
        logger.error(f"错误: 嵌入文件未找到 - {e}")
        logger.error("请确保 llm_enhancement/ 目录下的预处理脚本已成功运行，或检查 run.sh 的逻辑。")
        return  # 退出程序
    except Exception as e:
        logger.error(f"加载嵌入文件时发生未知错误: {e}", exc_info=True)
        return

    logger.info("语义画像嵌入准备完成。")

    # --- 4. 初始化模型 ---
    logger.info("--> Step 3: 初始化 CS_KGAT 模型...")
    try:
        # 【代码优化】: 将data对象直接传递给模型
        model = CS_KGAT(args, data, user_profiles_embed, item_profiles_embed)
        model.to(device)
        logger.info(model)
    except Exception as e:
        logger.error(f"模型初始化失败: {e}", exc_info=True)
        return
    logger.info("模型初始化完成。")

    # --- 5. 初始化并启动训练器 ---
    logger.info("--> Step 4: 初始化训练器并开始训练...")
    try:
        trainer = Trainer(args, model, data, device)
        trainer.train()
    except Exception as e:
        logger.error(f"训练过程中发生错误: {e}", exc_info=True)
        return
    logger.info("训练流程结束。")


if __name__ == '__main__':
    main()